package fmi23block1;

import java.io.FileReader;

import java.io.IOException;
import java.io.Reader;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.function.Function;

import org.jgrapht.nio.dimacs.DIMACSImporter;
import org.jgrapht.nio.ImportException;
import org.jgrapht.graph.DefaultDirectedGraph;
import org.jgrapht.graph.DefaultEdge;
import org.sat4j.minisat.SolverFactory;
import org.sat4j.specs.ContradictionException;
import org.sat4j.specs.IProblem;
import org.sat4j.specs.ISolver;
import org.sat4j.specs.IVec;
import org.sat4j.specs.IVecInt;
import org.sat4j.specs.TimeoutException;

/**
 * Main Class instantiating and running the SAT solver
 *
 */
public class ComputeColoring {

  private static final String STANDARD = "coloring";
  private static final String MAX = "maxColoring";

  /**
   * We deal with the command line parameters as:
   * 
   * 'coloring path': the first argument is the problem and the second argument is the path to the
   * file (for instance ./data/test1.graph), we compute a valid coloring
   * 
   * 'coloring path x y z': the first argument is the problem, the second argument is the path and
   * all the other arguments are the vertices that we want to color blue. For instance, 'coloring
   * ./data/test2.graph 1 2' computes a valid coloring that colors vertices 1 and 2 blue.
   * 
   * 'maxColoring path': the first argument is the problem, the second argument is the path. We
   * compute a blue-maximal valid coloring for the input graph.
   * 
   * @param args
   * @throws ImportException
   * @throws IOException
   * @throws TimeoutException
   */
  public static void main(String[] args) throws ImportException, IOException, TimeoutException {

    // Deal with the command line parameters
    if (args.length < 2) {
      System.out.println("USAGE: ComputeColoring problem path [blueVertex]*");
    } else if (args.length > 2 && args[0].equals(MAX)) {
      System.out.println("ERROR: You cannot specify a list of vertices for problem 'maxColoring'.");
    } else {
      String problem, path;
      List<Integer> blueVertices = new ArrayList<Integer>();

      problem = args[0];
      path = args[1];

      // Initialize graph
      DefaultDirectedGraph<Integer, DefaultEdge> graph;

      if (args.length >= 3) {
        for (int i = 2; i < args.length; i++) {
          blueVertices.add(Integer.parseInt(args[i]));
        }
      }
      switch (problem) {
        case STANDARD:

          // Read graph from File
          graph = readGraphFromFile(path);

          // Initialize Reduction
          SatEncodings reduction = new SatEncodings();
          IVec<IVecInt> cnf;

          // Compute the CNF formula
          if (args.length >= 3) {
            cnf = reduction.validColoring(graph, blueVertices);
            computeColoringSAT(graph, reduction, cnf);
          } else {
            cnf = reduction.validColoring(graph);
            computeColoringSAT(graph, reduction, cnf);
          }
          break;
        case MAX:
          // Read graph from File
          graph = readGraphFromFile(path);

          computeMaxColoring(graph);
          break;
        default:
          System.out.println("Task not supported: " + problem);
      }
    }
  }

  private static void computeColoringSAT(DefaultDirectedGraph<Integer, DefaultEdge> graph,
      SatEncodings reduction, IVec<IVecInt> cnf)
      throws TimeoutException, ImportException, IOException {
    // Output CNF
    System.out.println("CNF: " + clauseForm(cnf));

    // Initialize the SAT solver
    ISolver solver = SolverFactory.newDefault();

    try {
      // Feeds the formula to the solver
      solver.addAllClauses(cnf);

      IProblem instance = solver;

      solver.setTimeout(3600);

      // Computes a model / a coloring if possible
      if (instance.isSatisfiable()) {
        int model[] = instance.model();

        // Initialize Coloring
        Coloring coloring;

        // Compute a Coloring from the model
        coloring = reduction.getColoringFromModel(model);

        System.out.println("Coloring: " + coloring);
      } else {
        System.out.println("There is no valid coloring.");
      }
    } catch (ContradictionException e) {
      System.out.println("There is no valid coloring.");
    }

  }

  private static void computeMaxColoring(DefaultDirectedGraph<Integer, DefaultEdge> graph)
      throws TimeoutException, ImportException, IOException {
    SatEncodings reduction = new SatEncodings();
    int counter = 1;
    /*
     * Compute a first coloring
     */
    IVec<IVecInt> cnf = reduction.validColoring(graph);
    // System.out.println("CNF: "+ clauseForm(cnf));

    // Initialize empty coloring
    Coloring coloring = new Coloring();
    ISolver solver = SolverFactory.newDefault();

    try {
      // Feeds the formula to the solver
      solver.addAllClauses(cnf);

      IProblem instance = solver;

      solver.setTimeout(3600);

      if (instance.isSatisfiable()) {
        int model[] = instance.model();
        coloring = reduction.getColoringFromModel(model);
        System.out.println(counter + ": Coloring: " + coloring);
        counter++;
      } else {
        System.out.println("There is no valid coloring.");
      }
    } catch (ContradictionException e) {
      System.out.println("There is no valid coloring.");
    }
    /*
     * Iteratively enlarges the coloring
     */
    boolean solution_found = false;
    // Note that we need at most graph.vertexSet().size() iterations:
    // if the initial coloring colors 0 nodes blue, and we add one node
    // per iteration, we have at most graph.vertexSet().size() iterations.
    // (multiply * 10 just to be safe)
    for (int i = 0; i < graph.vertexSet().size() * 10; i++) {
      // CNF testing whether the given coloring is maximal
      cnf = reduction.maxColoring(graph, coloring);

      // System.out.println(counter + ": CNF: "+ clauseForm(cnf));
      solver = SolverFactory.newDefault();
      try {
        solver.addAllClauses(cnf);
      } catch (ContradictionException e) {
        // trivially unsatisfiable
        solution_found = true;
        break;
      }
      IProblem instance = solver;
      solver.setTimeout(3600);
      if (!instance.isSatisfiable()) {
        solution_found = true;
        break;
      } else {
        // updating the coloring
        int model[] = instance.model();
        coloring = reduction.getColoringFromModel(model);
        System.out.println(counter + ": Coloring: " + coloring);
        counter++;
      }
    }

    if (solution_found) {
      System.out.println("Max Coloring: " + coloring);
    } else {
      System.out.println(
          "You have reached the maximum number of iterations. Probably, your reduction is incorrect.");
    }
  }

  /**
   * Reads a directed graph in DIMACS format from file
   * 
   * @param path Path to the DIMACS file
   * @return DirectedGraph object corresponding to graph in the DIMACS file
   * @throws ImportException
   * @throws IOException
   */
  private static DefaultDirectedGraph<Integer, DefaultEdge> readGraphFromFile(String path)
      throws ImportException, IOException {

    DIMACSImporter<Integer, DefaultEdge> importer = new DIMACSImporter<Integer, DefaultEdge>();
    DefaultDirectedGraph<Integer, DefaultEdge> importDiGraph =
        new DefaultDirectedGraph<Integer, DefaultEdge>(DefaultEdge.class);
    Reader reader;
    reader = new FileReader(path);
    Function<Integer, Integer> vertexFactory = x -> x;
    importer.setVertexFactory(vertexFactory);
    importer.importGraph(importDiGraph, reader);
    reader.close();

    return importDiGraph;
  }

  /**
   * Get a nice String representation of CNFs
   * 
   * @param cnf
   * @return
   */
  public static String clauseForm(IVec<IVecInt> cnf) {
    StringBuffer clauseForm = new StringBuffer();
    clauseForm.append("{");

    Iterator<IVecInt> iterator = cnf.iterator();
    IVecInt clause;
    while (iterator.hasNext()) {
      clause = iterator.next();
      clauseForm.append("{");
      clauseForm.append(clause);
      clauseForm.append("}");
      if (iterator.hasNext()) {
        clauseForm.append(",");
      }
    }
    clauseForm.append("}");
    return clauseForm.toString();
  }

}

